# -*- coding: utf-8 -*-
"""
技术指标计算工具模块
提供各种常用技术指标的计算方法
"""
import pandas as pd
import numpy as np
import logging

logger = logging.getLogger(__name__)


class TechnicalIndicators:
    """技术指标计算工具类"""
    
    @staticmethod
    def calculate_ma(df: pd.DataFrame, periods: list = [5, 10, 20, 30, 60, 120]) -> pd.DataFrame:
        """
        计算移动平均线（Simple Moving Average）
        
        Args:
            df: 包含close_price列的DataFrame
            periods: 均线周期列表
            
        Returns:
            添加了均线列的DataFrame
        """
        if df.empty:
            return df
        
        for period in periods:
            if len(df) >= period:
                df[f'ma{period}'] = df['close_price'].rolling(window=period).mean()
        
        return df
    
    @staticmethod
    def calculate_ema(df: pd.DataFrame, column: str = 'close_price', period: int = 12) -> pd.Series:
        """
        计算指数移动平均线（Exponential Moving Average）
        
        Args:
            df: DataFrame
            column: 要计算EMA的列名
            period: 周期
            
        Returns:
            EMA序列
        """
        return df[column].ewm(span=period, adjust=False).mean()
    
    @staticmethod
    def calculate_macd(df: pd.DataFrame, fast_period: int = 12, 
                      slow_period: int = 26, signal_period: int = 9) -> pd.DataFrame:
        """
        计算MACD指标
        
        MACD (Moving Average Convergence Divergence)
        - DIF (快线) = EMA(12) - EMA(26)
        - DEA (慢线/信号线) = EMA(DIF, 9)
        - HIST (柱状图) = (DIF - DEA) * 2
        
        Args:
            df: 包含close_price列的DataFrame
            fast_period: 快线周期，默认12
            slow_period: 慢线周期，默认26
            signal_period: 信号线周期，默认9
            
        Returns:
            添加了macd_dif, macd_dea, macd_hist列的DataFrame
        """
        if df.empty or len(df) < slow_period:
            return df
        
        # 计算快速和慢速EMA
        ema_fast = df['close_price'].ewm(span=fast_period, adjust=False).mean()
        ema_slow = df['close_price'].ewm(span=slow_period, adjust=False).mean()
        
        # 计算DIF（快线）
        df['macd_dif'] = ema_fast - ema_slow
        
        # 计算DEA（慢线/信号线）
        df['macd_dea'] = df['macd_dif'].ewm(span=signal_period, adjust=False).mean()
        
        # 计算HIST（柱状图）
        df['macd_hist'] = (df['macd_dif'] - df['macd_dea']) * 2
        
        return df
    
    @staticmethod
    def calculate_rsi(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """
        计算RSI相对强弱指数（Relative Strength Index）
        
        RSI = 100 - (100 / (1 + RS))
        其中 RS = 平均上涨幅度 / 平均下跌幅度
        
        Args:
            df: 包含close_price列的DataFrame
            period: 周期，默认14
            
        Returns:
            添加了rsi列的DataFrame
        """
        if df.empty or len(df) < period:
            return df
        
        # 计算价格变化
        delta = df['close_price'].diff()
        
        # 分离上涨和下跌
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        # 计算平均涨幅和跌幅
        avg_gain = gain.rolling(window=period).mean()
        avg_loss = loss.rolling(window=period).mean()
        
        # 避免除零
        avg_loss = avg_loss.replace(0, 0.0001)
        
        # 计算RS和RSI
        rs = avg_gain / avg_loss
        df['rsi'] = 100 - (100 / (1 + rs))
        
        return df
    
    @staticmethod
    def calculate_kdj(df: pd.DataFrame, n: int = 9, m1: int = 3, m2: int = 3) -> pd.DataFrame:
        """
        计算KDJ指标（随机指标）
        
        Args:
            df: 包含high_price, low_price, close_price的DataFrame
            n: RSV周期，默认9
            m1: K值周期，默认3
            m2: D值周期，默认3
            
        Returns:
            添加了kdj_k, kdj_d, kdj_j列的DataFrame
        """
        if df.empty or len(df) < n:
            return df
        
        # 计算RSV
        low_min = df['low_price'].rolling(window=n).min()
        high_max = df['high_price'].rolling(window=n).max()
        
        # 避免除零
        rsv = 100 * (df['close_price'] - low_min) / (high_max - low_min + 0.0001)
        
        # 计算K值
        df['kdj_k'] = rsv.ewm(span=m1, adjust=False).mean()
        
        # 计算D值
        df['kdj_d'] = df['kdj_k'].ewm(span=m2, adjust=False).mean()
        
        # 计算J值
        df['kdj_j'] = 3 * df['kdj_k'] - 2 * df['kdj_d']
        
        return df
    
    @staticmethod
    def calculate_boll(df: pd.DataFrame, period: int = 20, std_dev: float = 2.0) -> pd.DataFrame:
        """
        计算布林带（Bollinger Bands）
        
        - 中轨 = MA(20)
        - 上轨 = 中轨 + 2 * 标准差
        - 下轨 = 中轨 - 2 * 标准差
        
        Args:
            df: 包含close_price列的DataFrame
            period: 周期，默认20
            std_dev: 标准差倍数，默认2.0
            
        Returns:
            添加了boll_mid, boll_upper, boll_lower列的DataFrame
        """
        if df.empty or len(df) < period:
            return df
        
        # 中轨
        df['boll_mid'] = df['close_price'].rolling(window=period).mean()
        
        # 标准差
        std = df['close_price'].rolling(window=period).std()
        
        # 上轨和下轨
        df['boll_upper'] = df['boll_mid'] + std_dev * std
        df['boll_lower'] = df['boll_mid'] - std_dev * std
        
        # 布林带宽度
        df['boll_width'] = (df['boll_upper'] - df['boll_lower']) / df['boll_mid']
        
        return df
    
    @staticmethod
    def calculate_atr(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """
        计算ATR平均真实波幅（Average True Range）
        
        TR = max(high - low, abs(high - prev_close), abs(low - prev_close))
        ATR = MA(TR, period)
        
        Args:
            df: 包含high_price, low_price, close_price的DataFrame
            period: 周期，默认14
            
        Returns:
            添加了atr列的DataFrame
        """
        if df.empty or len(df) < period:
            return df
        
        # 前一日收盘价
        prev_close = df['close_price'].shift(1)
        
        # 计算真实波幅
        tr1 = df['high_price'] - df['low_price']
        tr2 = abs(df['high_price'] - prev_close)
        tr3 = abs(df['low_price'] - prev_close)
        
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        
        # 计算ATR
        df['atr'] = tr.rolling(window=period).mean()
        
        return df
    
    @staticmethod
    def calculate_volume_indicators(df: pd.DataFrame) -> pd.DataFrame:
        """
        计算成交量相关指标
        
        - avg_vol_5: 5日平均成交量
        - avg_vol_10: 10日平均成交量
        - avg_vol_20: 20日平均成交量
        - vol_ratio: 量比（当日成交量 / 5日平均成交量）
        
        Args:
            df: 包含volume列的DataFrame
            
        Returns:
            添加了成交量指标列的DataFrame
        """
        if df.empty:
            return df
        
        # 计算不同周期的平均成交量
        if len(df) >= 5:
            df['avg_vol_5'] = df['volume'].rolling(window=5).mean()
        if len(df) >= 10:
            df['avg_vol_10'] = df['volume'].rolling(window=10).mean()
        if len(df) >= 20:
            df['avg_vol_20'] = df['volume'].rolling(window=20).mean()
        
        # 计算量比
        if 'avg_vol_5' in df.columns:
            df['vol_ratio'] = df['volume'] / df['avg_vol_5'].replace(0, np.nan)
        
        return df
    
    @staticmethod
    def calculate_week_data(df: pd.DataFrame) -> pd.DataFrame:
        """
        将日线数据转换为周线数据并计算周线MACD
        
        Args:
            df: 日线数据DataFrame
            
        Returns:
            添加了week_macd_dif, week_macd_dea列的DataFrame
        """
        if df.empty or len(df) < 26:
            df['week_macd_dif'] = np.nan
            df['week_macd_dea'] = np.nan
            return df
        
        try:
            # 复制数据并设置日期为索引
            df_week = df.copy()
            df_week.set_index('trade_date', inplace=True)
            
            # 转换为周线数据
            week_close = df_week['close_price'].resample('W').last()
            
            # 计算周线MACD
            ema_fast = week_close.ewm(span=12, adjust=False).mean()
            ema_slow = week_close.ewm(span=26, adjust=False).mean()
            week_dif = ema_fast - ema_slow
            week_dea = week_dif.ewm(span=9, adjust=False).mean()
            
            # 将周线数据映射回日线
            week_dif_daily = week_dif.reindex(df_week.index, method='ffill')
            week_dea_daily = week_dea.reindex(df_week.index, method='ffill')
            
            df['week_macd_dif'] = week_dif_daily.values
            df['week_macd_dea'] = week_dea_daily.values
            
        except Exception as e:
            logger.warning(f"计算周线MACD失败: {e}")
            df['week_macd_dif'] = np.nan
            df['week_macd_dea'] = np.nan
        
        return df
    
    @staticmethod
    def add_previous_values(df: pd.DataFrame, columns: list, periods: int = 1) -> pd.DataFrame:
        """
        添加指定列的前N期数值
        
        Args:
            df: DataFrame
            columns: 列名列表
            periods: 前N期，默认1（前一日）
            
        Returns:
            添加了prev_列的DataFrame
        """
        if df.empty:
            return df
        
        for col in columns:
            if col in df.columns:
                df[f'prev_{col}'] = df[col].shift(periods)
        
        return df
    
    @staticmethod
    def calculate_historical_extremes(df: pd.DataFrame, periods: list = [20, 60, 120]) -> pd.DataFrame:
        """
        计算历史极值（最高价、最低价等）
        
        Args:
            df: DataFrame
            periods: 周期列表
            
        Returns:
            添加了历史极值列的DataFrame
        """
        if df.empty:
            return df
        
        for period in periods:
            if len(df) >= period:
                # 最高价
                if 'high_price' in df.columns:
                    df[f'highest_high_{period}'] = df['high_price'].rolling(window=period).max()
                
                # 最低价
                if 'low_price' in df.columns:
                    df[f'lowest_low_{period}'] = df['low_price'].rolling(window=period).min()
                
                # 最高收盘价
                if 'close_price' in df.columns:
                    df[f'highest_close_{period}'] = df['close_price'].rolling(window=period).max()
                
                # 最低收盘价
                if 'close_price' in df.columns:
                    df[f'lowest_close_{period}'] = df['close_price'].rolling(window=period).min()
        
        return df
    
    @staticmethod
    def calculate_all_indicators(df: pd.DataFrame) -> pd.DataFrame:
        """
        计算所有常用技术指标（一站式计算）
        
        Args:
            df: 包含日线数据的DataFrame
            
        Returns:
            添加了所有技术指标的DataFrame
        """
        if df.empty:
            return df
        
        logger.info("开始计算技术指标...")
        
        # 移动平均线
        df = TechnicalIndicators.calculate_ma(df, periods=[5, 10, 20, 30, 60, 120])
        
        # MACD
        df = TechnicalIndicators.calculate_macd(df)
        
        # RSI
        df = TechnicalIndicators.calculate_rsi(df)
        
        # KDJ
        if 'high_price' in df.columns and 'low_price' in df.columns:
            df = TechnicalIndicators.calculate_kdj(df)
        
        # 布林带
        df = TechnicalIndicators.calculate_boll(df)
        
        # ATR
        if 'high_price' in df.columns and 'low_price' in df.columns:
            df = TechnicalIndicators.calculate_atr(df)
        
        # 成交量指标
        if 'volume' in df.columns:
            df = TechnicalIndicators.calculate_volume_indicators(df)
        
        # 周线MACD
        if 'trade_date' in df.columns:
            df = TechnicalIndicators.calculate_week_data(df)
        
        # 前一日数值
        df = TechnicalIndicators.add_previous_values(
            df, 
            columns=['close_price', 'macd_dif', 'macd_dea', 'macd_hist']
        )
        
        # 历史极值
        df = TechnicalIndicators.calculate_historical_extremes(df, periods=[20, 60, 120])
        
        # 添加MACD历史极值
        if 'macd_dif' in df.columns:
            for period in [20, 60, 120]:
                if len(df) >= period:
                    df[f'lowest_macd_dif_{period}'] = df['macd_dif'].rolling(window=period).min()
                    df[f'highest_macd_dif_{period}'] = df['macd_dif'].rolling(window=period).max()
        
        # 添加历史均线值
        for ma_period in [60, 120]:
            ma_col = f'ma{ma_period}'
            if ma_col in df.columns:
                df[f'{ma_col}_10d_ago'] = df[ma_col].shift(10)
                df[f'{ma_col}_20d_ago'] = df[ma_col].shift(20)
        
        logger.info("技术指标计算完成")
        
        return df


# 便捷函数
def calculate_macd(df: pd.DataFrame) -> pd.DataFrame:
    """便捷函数：计算MACD"""
    return TechnicalIndicators.calculate_macd(df)


def calculate_rsi(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
    """便捷函数：计算RSI"""
    return TechnicalIndicators.calculate_rsi(df, period)


def calculate_all_indicators(df: pd.DataFrame) -> pd.DataFrame:
    """便捷函数：计算所有指标"""
    return TechnicalIndicators.calculate_all_indicators(df)


